package main

import (
    "unsafe"
    "go-os-qemu/pkg/vga"
)

type Task struct {
    esp uintptr
    id  int
    name string
    state int // 0=ready,1=running,2=dead
    stack []byte
    arg   string
}

var (
    tasks      [8]*Task
    taskCount  int
    currentIdx int
    preemptEnabled bool
    quantumTicks   uint32 = 8
)

var (
    schedTicks uint32
)

// ctx_switch is implemented in boot/context.asm
//go:linkname ctx_switch ctx_switch
func ctx_switch(old *uintptr, new uintptr)

// Prepare initial stacks using ASM helpers to avoid func() -> uintptr conversion.
//go:linkname ctx_make_taskA ctx_make_taskA
func ctx_make_taskA(top uintptr) uintptr

//go:linkname ctx_make_taskB ctx_make_taskB
func ctx_make_taskB(top uintptr) uintptr

// Additional entry constructors from ASM
//go:linkname ctx_make_taskElf ctx_make_taskElf
func ctx_make_taskElf(top uintptr) uintptr

//go:linkname ctx_make_taskWorker ctx_make_taskWorker
func ctx_make_taskWorker(top uintptr) uintptr

func initScheduler() {
    // Initialize empty task table; demo tasks created via createDemoTasks()
    taskCount = 0
    currentIdx = 0
    preemptEnabled = true
    vga.Println("[sched] runtime ready")
}

// yield cooperatively switches to the next task.
func yield() {
    if taskCount == 0 { return }
    // Find next runnable
    start := (currentIdx + 1) % taskCount
    idx := start
    for {
        if tasks[idx] != nil && tasks[idx].state == 0 {
            cur := tasks[currentIdx]
            n := tasks[idx]
            tasks[currentIdx].state = 0
            n.state = 1
            ctx_switch(&cur.esp, n.esp)
            currentIdx = idx
            return
        }
        idx = (idx + 1) % taskCount
        if idx == start { break }
    }
}

// preempt switches to the next ready task; intended for IRQ context.
func preempt() {
    if taskCount == 0 { return }
    start := (currentIdx + 1) % taskCount
    idx := start
    for {
        if tasks[idx] != nil && tasks[idx].state == 0 {
            cur := tasks[currentIdx]
            n := tasks[idx]
            tasks[currentIdx].state = 0
            n.state = 1
            ctx_switch(&cur.esp, n.esp)
            currentIdx = idx
            return
        }
        idx = (idx + 1) % taskCount
        if idx == start { break }
    }
}

// schedulerTick is called from the timer IRQ to implement time slicing.
func schedulerTick() {
    if !preemptEnabled { return }
    schedTicks++
    if taskCount > 0 && (schedTicks%quantumTicks) == 0 {
        preempt()
    }
}

// Runtime controls
func setQuantum(n int) {
    if n < 1 { n = 1 }
    if n > 1000 { n = 1000 }
    quantumTicks = uint32(n)
    vga.SetColor(vga.ColorLightGray, vga.ColorBlack)
    vga.Print("[sched] quantum=")
    printDec(n)
    vga.Println("")
}

func setPreempt(on bool) {
    preemptEnabled = on
    vga.SetColor(vga.ColorLightGray, vga.ColorBlack)
    if on {
        vga.Println("[sched] preempt=on")
    } else {
        vga.Println("[sched] preempt=off")
    }
}

func startFirstTask() {
    if taskCount == 0 { return }
    // Switch from kernel context into first task
    next := tasks[0]
    next.state = 1
    var kernelESP uintptr
    ctx_switch(&kernelESP, next.esp)
    // After tasks eventually switch back, execution will resume here.
}

func taskA() {
    for {
        vga.SetColor(vga.ColorLightGreen, vga.ColorBlack)
        vga.Print("A")
        vga.SetColor(vga.ColorLightGray, vga.ColorBlack)
        yield()
    }
}

func taskB() {
    for {
        vga.SetColor(vga.ColorLightRed, vga.ColorBlack)
        vga.Print("B")
        vga.SetColor(vga.ColorLightGray, vga.ColorBlack)
        yield()
    }
}

// Exported entry wrappers for ASM to reference as initial return targets.
//go:export taskA_entry
func taskA_entry() { taskA() }

//go:export taskB_entry
func taskB_entry() { taskB() }

// Worker and ELF task entries
//go:export worker_entry
func worker_entry() { workerTask() }

//go:export elfTask_entry
func elfTask_entry() { elfTask() }

func workerTask() {
    for {
        t := tasks[currentIdx]
        if t == nil { return }
        vga.SetColor(vga.ColorLightGray, vga.ColorBlack)
        if len(t.name) > 0 {
            vga.Print(string(t.name[0]))
        } else {
            vga.Print("?")
        }
        yield()
    }
}

func elfTask() {
    t := tasks[currentIdx]
    if t == nil { return }
    // If no path, mark dead
    if t.arg == "" {
        t.state = 2
        yield()
        return
    }
    // Resolve and load, then run
    full := resolvePath(t.arg)
    data, ok := rootFS.ReadFile(full)
    if !ok {
        vga.SetColor(vga.ColorLightRed, vga.ColorBlack)
        vga.Println("[elf] no such file: ")
        vga.Println(full)
        t.state = 2
        yield()
        return
    }
    runBinary(data)
    // After program returns, mark dead and yield away
    t.state = 2
    vga.SetColor(vga.ColorLightGray, vga.ColorBlack)
    vga.Print("[elf] exit pid=")
    printDec(t.id)
    vga.Println("")
    yield()
}

// Public runtime API
// createDemoTasks initializes two demo tasks to showcase runtime.
func createDemoTasks() {
    if taskCount > 0 { return }
    stackA := make([]byte, 4096)
    stackB := make([]byte, 4096)
    topA := uintptr(unsafe.Pointer(&stackA[len(stackA)-1])) &^ uintptr(0xF)
    topB := uintptr(unsafe.Pointer(&stackB[len(stackB)-1])) &^ uintptr(0xF)
    tasks[0] = &Task{esp: ctx_make_taskA(topA), id: 0, name: "taskA", state: 0, stack: stackA}
    tasks[1] = &Task{esp: ctx_make_taskB(topB), id: 1, name: "taskB", state: 0, stack: stackB}
    taskCount = 2
    currentIdx = 0
    vga.Println("[sched] demo tasks created")
}

// newWorker spawns a simple worker task that prints its name initial.
func newWorker(name string) int {
    // find free slot
    idx := -1
    for i := 0; i < len(tasks); i++ {
        if tasks[i] == nil { idx = i; break }
    }
    if idx == -1 { vga.Println("[sched] no free task slot"); return -1 }
    stack := make([]byte, 4096)
    top := uintptr(unsafe.Pointer(&stack[len(stack)-1])) &^ uintptr(0xF)
    esp := ctx_make_taskWorker(top)
    tasks[idx] = &Task{esp: esp, id: idx, name: name, state: 0, stack: stack}
    if idx >= taskCount { taskCount = idx + 1 }
    vga.SetColor(vga.ColorLightGray, vga.ColorBlack)
    vga.Print("[sched] new pid=")
    printDec(idx)
    vga.Print(" name=")
    vga.Println(name)
    return idx
}

// spawnElfTask starts an ELF program as a task.
func spawnElfTask(path string) int {
    idx := -1
    for i := 0; i < len(tasks); i++ {
        if tasks[i] == nil { idx = i; break }
    }
    if idx == -1 { vga.Println("[sched] no free task slot"); return -1 }
    stack := make([]byte, 8192)
    top := uintptr(unsafe.Pointer(&stack[len(stack)-1])) &^ uintptr(0xF)
    esp := ctx_make_taskElf(top)
    name := "elf"
    tasks[idx] = &Task{esp: esp, id: idx, name: name, state: 0, stack: stack, arg: path}
    if idx >= taskCount { taskCount = idx + 1 }
    vga.SetColor(vga.ColorLightGray, vga.ColorBlack)
    vga.Print("[sched] spawn ELF pid=")
    printDec(idx)
    vga.Print(" path=")
    vga.Println(path)
    return idx
}

// killTask marks a task dead and clears its slot. If killing current, preempt.
func killTask(pid int) {
    if pid < 0 || pid >= taskCount { vga.Println("[sched] bad pid"); return }
    t := tasks[pid]
    if t == nil { vga.Println("[sched] no such pid"); return }
    t.state = 2
    tasks[pid] = nil
    t.stack = nil
    vga.SetColor(vga.ColorLightGray, vga.ColorBlack)
    vga.Print("[sched] killed pid=")
    printDec(pid)
    vga.Println("")
    if pid == currentIdx {
        preempt()
    }
}

// ps prints current tasks
func ps() {
    vga.SetColor(vga.ColorLightGray, vga.ColorBlack)
    vga.Println("PID  NAME   STATE")
    for i := 0; i < taskCount; i++ {
        if tasks[i] == nil { continue }
        printDec(tasks[i].id)
        vga.Print("   ")
        vga.Print(tasks[i].name)
        vga.Print("    ")
        switch tasks[i].state {
        case 0: vga.Println("ready")
        case 1: vga.Println("running")
        case 2: vga.Println("dead")
        default: vga.Println("?")
        }
    }
}

// anyTaskRunning reports whether any task is currently marked running.
func anyTaskRunning() bool {
    for i := 0; i < taskCount; i++ {
        if tasks[i] != nil && tasks[i].state == 1 { return true }
    }
    return false
}